import os
import logging

# 配置日志
logger = logging.getLogger(__name__)

class GarbageLabels:
    def __init__(self):
        # 预测类别（214类）到详细信息的映射
        self.label_to_category = {}
        
        # 主类别（4类）映射
        self.main_categories = {
            0: "可回收物",
            1: "有害垃圾",
            2: "厨余垃圾",
            3: "其他垃圾"
        }
        
        # 垃圾分类规则字典
        self.garbage_rules = {
            '厨余垃圾': {
                'definition': '厨余垃圾，也称为湿垃圾，是指居民日常生活及食品加工、饮食服务、单位供餐等活动中产生的垃圾，包括剩菜剩饭、骨头、菜根菜叶、果皮等食品类废物。',
                'rules': '厨余垃圾应从产生时就与其他品种垃圾分开收集，投放前沥干水分。有包装物的湿垃圾应将包装物去除后分类投放，包装物应投放至对应的可回收物或干垃圾收集容器中。盛放湿垃圾的废弃塑料袋，应投放至干垃圾收集容器中。'
            },
            '其他垃圾': {
                'definition': '其他垃圾，也称为干垃圾，是指除了厨余垃圾、有害垃圾、可回收物以外的生活垃圾，即现环卫体系主要收集和处理的垃圾。其他垃圾的特点是危害较小，但无再次利用的价值。',
                'rules': '干垃圾应投入干垃圾收集容器，并保持周边环境整洁。生活中遇到成分复杂、不易分离归集的物品，建议作为其他垃圾处理。用过的餐巾纸、尿片宜作为其他垃圾处理。'
            },
            '可回收物': {
                'definition': '可回收物是指回收后经过再加工可以成为生产原料或者经过整理可以再利用的物品，主要包括废纸类、塑料类、玻璃类、金属类、电子废弃物类、织物类等。',
                'rules': '分类投放可回收物时，应保持清洁干燥，避免污染。其中：废纸应保持平整；立体包装物应清空内容物，清洁后压扁投放；废玻璃制品应轻放，有尖锐边角的应包裹后投放。'
            },
            '有害垃圾': {
                'definition': '有害垃圾是指含有害物质，需要特殊安全处理的生活垃圾。主要包括过期药品、水银温度计、废血压计、废杀虫剂及包装物、消毒剂及包装物、废油漆和溶剂及包装物、废矿物油及包装物、废荧光灯管（包括日光灯、节能灯，不包括LED灯管）、废充电电池、废纽扣电池、废镍镉电池、铅蓄电池和氧化汞电池以及电子类危险废物、家用化学品等。',
                'rules': '分类投放有害垃圾时，应注意轻放。其中：废灯管等易破损的有害垃圾应连带包装或包裹后投放；废弃药品宜连带包装一并投放；杀虫剂等压力罐装容器，应排空内容物后投放。'
            }
        }

    def load_label_mappings(self, label_file='dir_label.txt'):
        """加载标签映射"""
        logger.info('正在加载标签映射...')
        try:
            with open(label_file, 'r', encoding='utf-8') as f:
                for line in f:
                    parts = line.strip().split('\t')
                    if len(parts) >= 3:  # 确保有3个部分：类别名、索引、主类别索引
                        category_name = parts[0]  # 完整类别名称
                        label_id = int(parts[1])  # 类别索引
                        main_class_id = int(parts[2])  # 主类别索引
                        
                        # 从类别名称中提取主类别和详细类别
                        name_parts = category_name.split('_', 1)
                        if len(name_parts) == 2:
                            main_class = name_parts[0]
                            detailed_class = name_parts[1]
                        else:
                            main_class = category_name
                            detailed_class = category_name
                        
                        self.label_to_category[label_id] = {
                            'name': category_name,
                            'main_class': main_class,
                            'detailed_class': detailed_class
                        }
            logger.info(f'成功加载 {len(self.label_to_category)} 个标签')
            return True
        except Exception as e:
            logger.error(f'加载标签映射出错: {str(e)}')
            return False
            
    def load_answers(self, answer_file='answer.txt'):
        """加载答案文件中的垃圾分类数据"""
        answers = []
        try:
            with open(answer_file, 'r', encoding='utf-8') as f:
                for line in f:
                    # 移除开头的 ├─ 和空格，以及结尾的空白
                    category = line.strip().replace('├─ ', '').strip()
                    if category:
                        answers.append(category)
            return answers
        except Exception as e:
            logger.error(f"加载答案文件失败: {str(e)}")
            return []

    def get_prediction_info(self, predicted_idx):
        """获取预测的类别信息"""
        try:
            # 使用之前加载的label_to_category获取详细信息
            category_info = self.label_to_category.get(predicted_idx, None)
            if category_info:
                main_class = category_info['main_class']
                detailed_class = category_info['detailed_class']
                rules = self.garbage_rules.get(main_class, {
                    'definition': f'{main_class}类垃圾',
                    'rules': f'请按照{main_class}的相关规定处理'
                })
                return {
                    'main_class': main_class,
                    'detailed_class': detailed_class,
                    'definition': rules['definition'],
                    'rules': rules['rules']
                }
            
            # 如果找不到详细信息，使用主类别
            main_class = self.main_categories.get(predicted_idx // 100, '未知')  # 使用整除获取主类别索引
            rules = self.garbage_rules.get(main_class, {
                'definition': f'{main_class}类垃圾',
                'rules': f'请按照{main_class}的相关规定处理'
            })
            return {
                'main_class': main_class,
                'detailed_class': main_class,  # 如果没有详细类别，使用主类别作为详细类别
                'definition': rules['definition'],
                'rules': rules['rules']
            }
        except Exception as e:
            logger.error(f'获取预测信息出错: {str(e)}')
            # 返回默认值而不是None
            return {
                'main_class': '未知',
                'detailed_class': '未知',
                'definition': '未知类型垃圾',
                'rules': '请按照相关规定处理'
            }

    def get_class_from_path(self, image_path, dataset_path):
        """从图片路径中提取类别信息"""
        try:
            # 获取相对路径
            relative_path = os.path.relpath(image_path, dataset_path)
            # 获取父目录名称，其中包含类别信息
            parent_dir = os.path.dirname(relative_path)
            
            # 从目录名称中提取主类别和详细类别
            parts = parent_dir.split('_', 1)
            if len(parts) == 2:
                main_category, class_name = parts
            else:
                main_category = class_name = parts[0]
                
            return {
                "main_category": main_category,
                "class_name": class_name
            }
        except Exception as e:
            logger.error(f'从路径提取类别信息失败 {image_path}: {str(e)}')
            return None

    def get_search_suggestions(self):
        """获取搜索提示词列表"""
        return [
            # 食物分类
            "水果", "蔬菜", "零食", "坚果", "蛋白质", 
            # 材质分类
            "玻璃", "塑料", "金属", "纸张",  "电子产品",
            # 有害垃圾
            "厨房", "办公", "餐饮", "客厅电器", "书房文具", 
            # 其他垃圾
            "危险", "药片", "化学用品", "电子电器", "白色垃圾",
            # 垃圾分类
            "厨余垃圾", "有害垃圾", "可回收物", "其他垃圾"
        ]